# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags
from concurrent.futures import ProcessPoolExecutor

import collections
import functools
import logging

import numpy
import pandas
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from coin.base.datetime_util import to_timestamp_int

from coin.base.book.types import TradeSide
import coin.strategy.mm.tool.archive_base as abase


class TbComparison(object):
  def __init__(self, products, start_ts=None, end_ts=None):
    self._products = products
    self._tick_sizes = (0.5, 0.01)
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._end_ts = to_timestamp_int(end_ts or 0)
    self._fields_out = []
    self._fields_view = []
    self._time_set = set()
    self._agg_book = {}
    self._tsigned_cum_qty = collections.defaultdict(float)
    self._trades = {product: collections.OrderedDict() for product in products}
    self._bbos = {product: collections.OrderedDict() for product in products}
    self._instrs = {product: collections.OrderedDict() for product in products}
    self._last_buy_trade = collections.defaultdict(lambda: (None, None))
    self._last_sell_trade = collections.defaultdict(lambda: (None, None))

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      if not hasattr(book_builder, "ProductType") or isinstance(product, book_builder.ProductType):
        book_builder.subscribe(product, functools.partial(self.on_book, product))
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def on_trade(self, product, trade):
    if trade.timestamp < self._start_ts:
      return

    if self._end_ts and self._end_ts < trade.timestamp:
      return

    ts_in_unit = int(trade.timestamp * 1e-7) * 1e7

    tradesign = 1 if trade.side == TradeSide.TRADE_BUY_SIDE else -1

    self._tsigned_cum_qty[product] += trade.qty * tradesign

    if tradesign == 1:
      self._last_buy_trade[product] = (trade.price, trade.qty)
    else:
      self._last_sell_trade[product] = (trade.price, trade.qty)

    self._trades[product][ts_in_unit] = {
        'ts_in_unit': ts_in_unit,
        'timestamp': trade.timestamp,
        'trade_timestamp': trade.timestamp,
        'type': 'T',
        'last_buy_trade_price': self._last_buy_trade[product][0],
        'last_buy_trade_qty': self._last_buy_trade[product][1],
        'last_sell_trade_price': self._last_sell_trade[product][0],
        'last_sell_trade_qty': self._last_sell_trade[product][1],
        'tradecumsq': self._tsigned_cum_qty[product]
    }

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < book.timestamp:
      return

    try:
      bid0_price = book.bid0().price
      ask0_price = book.ask0().price
    except (AttributeError, IndexError):
      return

    ts_in_unit = int(book.timestamp * 1e-7) * 1e7

    self._bbos[product][ts_in_unit] = {
        'ts_in_unit': ts_in_unit,
        'type': 'B',
        'timestamp': book.timestamp,
        'book_timestamp': book.timestamp,
        'bid0p': bid0_price,
        'ask0p': ask0_price
    }

  def analyze(self, out_dir, trading_date_str, machine):
    dfs = []
    for product in self._products:
      bdf = pandas.DataFrame(self._bbos[product]).T
      offset_err = collections.OrderedDict()
      # for offset in [-2e9, -500e6, -100e6, 0, 100e6, 500e6, 2e9]:
      for offset in [0]:
        tdf = pandas.DataFrame(self._trades[product]).T
        tdf['timestamp'] += offset
        mdf = pandas.concat([tdf.reset_index(drop=True), bdf.reset_index(drop=True)],
                            sort=False).sort_values('timestamp').fillna(method='ffill')

        tmdf = mdf.loc[mdf.type == 'T'].reset_index(drop=True)
        desc = (tmdf['trade_timestamp'] - tmdf['book_timestamp']).describe()
        ddf = desc.reset_index()
        ddf['date'] = trading_date_str
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        tds = tmdf['trade_timestamp'] - tmdf['book_timestamp']
        tds = tds[~numpy.isnan(tds)]
        sortedtds = numpy.array(sorted(tds[tds < numpy.quantile(tds, 0.9)] * 1e-6))

        tds2 = numpy.array(tdf['trade_timestamp'].diff()).astype(float)
        tds2 = tds2[~numpy.isnan(tds2)]
        sortedtds2 = numpy.array(sorted(tds2[tds2 < 1200e6] * 1e-6))

        dfs.append({
            'trading_date_str': trading_date_str,
            'T timestamp - B snapshot timestamp': sortedtds.mean()
        })
        plt.subplot(211)
        plt.plot(sortedtds, lw=0.5, markersize=0.5)
        plt.title(trading_date_str)
        plt.ylabel("trade timestamp - book timestamp dist.")

        plt.subplot(212)
        plt.plot(sortedtds2, lw=0.5, markersize=0.5)
        plt.savefig(f"{out_dir}/{trading_date_str}.png")
        plt.close()
        '''
        import pdb; pdb.set_trace()
        desc.to_dataframe()

        mdf = mdf.loc[mdf.type == 'B'].reset_index(drop=True)
        offset_err[offset] = {
            'trading_date': trading_date_str,
            'stdev': (mdf.ask0p - mdf.last_buy_trade_price).std() + (
                mdf.bid0p - mdf.last_sell_trade_price).std(),
            'predicts': (mdf.last_buy_trade_price - mdf.ask0p).mean() + (
                mdf.bid0p - mdf.last_sell_trade_price).mean(),
            'predicts_dist': (mdf.last_buy_trade_price > mdf.ask0p).mean() + (
                mdf.bid0p > mdf.last_sell_trade_price).mean()
            }
            '''
        # dfs.append(ddf)

      # dfs.append(pandas.DataFrame(offset_err).T)
    return dfs


def launch(baseflags, fsflags, symbols, subsymbols, out_dir):
  plt.rcParams['figure.figsize'] = 12, 8
  plt.rcParams['font.size'] = 10
  plt.rcParams['legend.fontsize'] = 10
  plt.rcParams['xtick.labelsize'] = 10
  plt.rcParams['ytick.labelsize'] = 10

  products = []
  if symbols is not None:
    products.extend(abase.symbols_to_products(symbols, baseflags))

  if subsymbols is not None:
    products.extend(abase.subsymbols_to_products(subsymbols, baseflags))

  start_time, end_time = abase.get_time_range(baseflags)

  strategy = TbComparison(products=products, start_ts=start_time, end_ts=end_time)

  abase.run_from_archive_base(baseflags=baseflags,
                              feed_sub_flags_proto=fsflags,
                              products=products,
                              on_book_reset=strategy.on_book_reset,
                              api_version_override={
                                  ('Futures', 'Okex'): 'v3',
                                  ('Spot', 'Okex'): 'v3',
                              })

  dfs = strategy.analyze(out_dir,
                         trading_date_str=baseflags.trading_date,
                         machine=fsflags.feed_machine)
  return dfs


def launch_date(baseflags, fsflags, out_dir, symbols, subsymbols):
  return launch(baseflags, fsflags, symbols, subsymbols, out_dir)


def main(argv):
  subsymbols = None if flags.FLAGS.subsymbols is None else flags.FLAGS.subsymbols.split(",")

  jobs = []

  for trading_date in abase.get_trading_dates(flags.FLAGS.trading_date):
    flags.FLAGS.trading_date = trading_date.strftime("%Y%m%d")
    baseflags = abase.get_base_flags()
    fsflags = abase.get_feed_subscription_flags()
    jobs.append((launch_date, baseflags, fsflags, flags.FLAGS.out_dir, [], subsymbols))

  if len(jobs) == 1:
    print(jobs[0][0](*jobs[0][1:]))
  else:
    futures = []
    dfs = []
    with ProcessPoolExecutor(max_workers=8) as executor:
      for job in jobs:
        futures.append(executor.submit(job[0], *job[1:]))
      for future in futures:
        dfs.extend(future.result())
    for df in dfs:
      print(df)


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  # fill up default settings
  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'
  flags.DEFINE_string('postfix', None, 'filename postfix')
  flags.DEFINE_string('subsymbols', None, 'subscript symbol name')
  flags.DEFINE_string('out_dir', 'tb', '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')

  app.run(main)
